import os
import cv2
import argparse

def augment_image(path, file, output):
    image = cv2.imread(path)
    heigh = image.shape[0]
    width = image.shape[1]
    filename = file.split(".")
    center = (width // 2, heigh // 2)

    # [0, 1, 2, 3] * 90 , No rotation, 90 counterclockwise, 180 counterclockwise, 270 counterclockwise
    for i in range(4):
        change_image = cv2.warpAffine(image, cv2.getRotationMatrix2D(center, i*90, 1.0), (width, heigh))
        cv2.imwrite("{}/{}_{}.{}".format(output, filename[0], str(i+1), filename[1]), change_image)
    
    change_image = cv2.flip(image, 0)# 0 flipped vertically, 1 flipped horizontally
    cv2.imwrite("{}/{}_{}.{}".format(output, filename[0], "5", filename[1]), change_image)

    change_image = cv2.flip(image, 1)# 0 flipped vertically, 1 flipped horizontally
    cv2.imwrite("{}/{}_{}.{}".format(output, filename[0], "6", filename[1]), change_image)

if __name__=="__main__":
    parser = argparse.ArgumentParser(description='Split image')
    parser.add_argument('--input', '-i', type=str, required=True, help='Path to the input image')
    parser.add_argument('--output', '-o', type=str, required=True, help='Path to the output image')
    args = parser.parse_args()

    for root, dirs, files in os.walk(args.input):
        for file in files:
            if 'tif' in file:
                augment_image("{}/{}".format(root,file), file, args.output)
